package rabbit.flt.plugins.spring.r2dbc.low.plugin;

import org.springframework.data.r2dbc.core.FetchSpec;
import rabbit.flt.plugins.common.R2dbcSupportPlugin;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;
import java.util.function.Supplier;

public class DefaultDatabaseClientPlugin extends R2dbcSupportPlugin {

    @Override
    public Object after(Object objectEnhanced, Method method, Object[] args, Object result) {
        super.after(objectEnhanced, method, args, result);
        if (isTraceOpened()) {
            Supplier<String> sqlSupplier = getSqlSupplier(method, args);
            String sql = sqlSupplier.get().trim();
            FetchSpec spec = (FetchSpec) result;
            return new FetchSpec() {

                @Override
                public Mono<Integer> rowsUpdated() {
                    return (Mono<Integer>) wrapResult(spec.rowsUpdated(), method, sql);
                }

                @Override
                public Mono one() {
                    return wrapResult(spec.one(), method, sql);
                }

                @Override
                public Mono first() {
                    return wrapResult(spec.first(), method, sql);
                }

                @Override
                public Flux all() {
                    return wrapResult(spec.all(), method, sql);
                }
            };
        }
        return result;
    }

    protected Supplier<String> getSqlSupplier(Method method, Object[] args) {
        if (method.getName().equals("exchangeInsert")) {
            return (Supplier<String>) args[1];
        }
        return (Supplier<String>) args[0];
    }

}
